// Copyright 2010-2025 Google LLC
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Simple framework for choosing and distributing a solver "sub-tasks" on a set
// of threads.

#ifndef OR_TOOLS_SAT_SUBSOLVER_H_
#define OR_TOOLS_SAT_SUBSOLVER_H_

#include <algorithm>
#include <cmath>
#include <cstdint>
#include <functional>
#include <memory>
#include <string>
#include <utility>
#include <vector>

#include "absl/strings/string_view.h"
#include "ortools/sat/util.h"
#include "ortools/util/stats.h"

#if !defined(__PORTABLE_PLATFORM__)
#endif  // __PORTABLE_PLATFORM__

namespace operations_research {
namespace sat {

// The API used for distributing work. Each subsolver can generate tasks and
// synchronize itself with the rest of the world.
//
// Note that currently only the main thread interact with subsolvers. Only the
// tasks generated by GenerateTask() are executed in parallel in a threadpool.
class SubSolver {
 public:
  enum SubsolverType { FULL_PROBLEM, FIRST_SOLUTION, INCOMPLETE, HELPER };

  SubSolver(absl::string_view name, SubsolverType type)
      : name_(name), type_(type) {}
  virtual ~SubSolver() = default;

  // Synchronizes with the external world from this SubSolver point of view.
  // Also incorporate the results of the latest completed tasks if any.
  //
  // Note(user): The intended implementation for determinism is that tasks
  // update asynchronously (and so non-deterministically) global "shared"
  // classes, but this global state is incorporated by the Subsolver only when
  // Synchronize() is called.
  //
  // This is only called by the main thread in Subsolver creation order.
  virtual void Synchronize() = 0;

  // Returns true if this SubSolver is done and its memory can be freed. Note
  // that the *Loop(subsolvers) functions below takes a reference in order to be
  // able to clear the memory of a SubSolver as soon as it is done. Once this is
  // true, the subsolver in question will be deleted and never used again.
  //
  // This is needed since some subsolve can be done before the overal Solve() is
  // finished. This is the case for first solution subsolvers for instances.
  //
  // This is only called by the main thread in a sequential fashion.
  // Important: This is only called when there is currently no task from that
  // SubSolver in flight.
  virtual bool IsDone() { return false; }

  // Returns true iff GenerateTask() can be called.
  // This is only called by the main thread in a sequential fashion.
  virtual bool TaskIsAvailable() = 0;

  // Returns a task to run. The task_id is just an ever increasing counter that
  // correspond to the number of total calls to GenerateTask().
  //
  // TODO(user): We could use a more complex selection logic and pass in the
  // deterministic time limit this subtask should run for. Unclear at this
  // stage.
  //
  // This is only called by the main thread.
  virtual std::function<void()> GenerateTask(int64_t task_id) = 0;

  // Returns the total deterministic time spend by the completed tasks before
  // the last Synchronize() call.
  double deterministic_time() const { return deterministic_time_; }

  // Returns the name of this SubSolver. Used in logs.
  std::string name() const { return name_; }

  // Returns the type of the subsolver.
  SubsolverType type() const { return type_; }

  // Note that this is protected by the global execution mutex and so it is
  // called sequentially. Subclasses do not need to call this.
  void AddTaskDuration(double duration_in_seconds) {
    ++num_finished_tasks_;
    wall_time_ += duration_in_seconds;
    timing_.AddTimeInSec(duration_in_seconds);
  }

  // Note that this is protected by the global execution mutex and so it is
  // called sequentially. Subclasses do not need to call this.
  void NotifySelection() { ++num_scheduled_tasks_; }

  // This one need to be called by the Subclasses. Usually from Synchronize(),
  // or from the task itself it we execute a single task at the same time.
  void AddTaskDeterministicDuration(double deterministic_duration) {
    if (deterministic_duration <= 0) return;
    deterministic_time_ += deterministic_duration;
    dtiming_.AddTimeInSec(deterministic_duration);
  }

  std::string TimingInfo() const {
    // TODO(user): remove trailing "\n" from ValueAsString() or just build the
    // table line directly.
    std::string data = timing_.ValueAsString();
    if (!data.empty()) data.pop_back();
    return data;
  }

  std::string DeterministicTimingInfo() const {
    // TODO(user): remove trailing "\n" from ValueAsString().
    std::string data = dtiming_.ValueAsString();
    if (!data.empty()) data.pop_back();
    return data;
  }

  // Returns a score used to compare which tasks to schedule next.
  // We will schedule the LOWER score.
  //
  // Tricky: Note that this will only be called sequentially. The deterministic
  // time should only be used with the DeterministicLoop() because otherwise it
  // can be updated at the same time as this is called.
  double GetSelectionScore(bool deterministic) const {
    const double time = deterministic ? deterministic_time_ : wall_time_;
    const double divisor = num_scheduled_tasks_ > 0
                               ? static_cast<double>(num_scheduled_tasks_)
                               : 1.0;

    // If we have little data, we strongly limit the number of task in flight.
    // This is needed if some LNS are stuck for a long time to not just only
    // schedule this type at the beginning.
    const int64_t in_flight = num_scheduled_tasks_ - num_finished_tasks_;
    const double confidence_factor =
        num_finished_tasks_ > 10 ? 1.0 : std::exp(in_flight);

    // We assume a "minimum time per task" which will be our base etimation for
    // the average running time of this task.
    return num_scheduled_tasks_ * std::max(0.1, time / divisor) *
           confidence_factor;
  }

 private:
  const std::string name_;
  const SubsolverType type_;

  int64_t num_scheduled_tasks_ = 0;
  int64_t num_finished_tasks_ = 0;

  // Sum of wall_time / deterministic_time.
  double wall_time_ = 0.0;
  double deterministic_time_ = 0.0;

  TimeDistribution timing_ = TimeDistribution("task time");
  TimeDistribution dtiming_ = TimeDistribution("task dtime");
};

// A simple wrapper to add a synchronization point in the list of subsolvers.
class SynchronizationPoint : public SubSolver {
 public:
  explicit SynchronizationPoint(absl::string_view name, std::function<void()> f)
      : SubSolver(name, HELPER), f_(std::move(f)) {}
  bool TaskIsAvailable() final { return false; }
  std::function<void()> GenerateTask(int64_t /*task_id*/) final {
    return nullptr;
  }
  void Synchronize() final { f_(); }

 private:
  std::function<void()> f_;
};

// Executes the following loop:
// 1/ Synchronize all in given order.
// 2/ generate and schedule one task from the current "best" subsolver.
// 3/ repeat until no extra task can be generated and all tasks are done.
//
// The complexity of each selection is in O(num_subsolvers), but that should
// be okay given that we don't expect more than 100 such subsolvers.
//
// Note that it is okay to incorporate "special" subsolver that never produce
// any tasks. This can be used to synchronize classes used by many subsolvers
// just once for instance.
void NonDeterministicLoop(std::vector<std::unique_ptr<SubSolver>>& subsolvers,
                          int num_threads, ModelSharedTimeLimit* time_limit);

// Similar to NonDeterministicLoop() except this should result in a
// deterministic solver provided that all SubSolver respect the Synchronize()
// contract.
//
// Executes the following loop:
// 1/ Synchronize all in given order.
// 2/ generate and schedule up to batch_size tasks using an heuristic to select
//    which one to run.
// 3/ wait for all task to finish.
// 4/ repeat until no task can be generated in step 2.
//
// If max_num_batches is > 0, stop after that many batches.
void DeterministicLoop(std::vector<std::unique_ptr<SubSolver>>& subsolvers,
                       int num_threads, int batch_size,
                       int max_num_batches = 0);

// Same as above, but specialized implementation for the case num_threads=1.
// This avoids using a Threadpool altogether. It should have the same behavior
// than the functions above with num_threads=1 and batch_size=1. Note that an
// higher batch size will not behave in the same way, even if num_threads=1.
void SequentialLoop(std::vector<std::unique_ptr<SubSolver>>& subsolvers);

}  // namespace sat
}  // namespace operations_research

#endif  // OR_TOOLS_SAT_SUBSOLVER_H_
